import numpy as np
import matplotlib.pyplot as plt
from PIL import Image

import sys
import os.path as osp
sys.path.append(osp.dirname(osp.dirname(osp.abspath(__file__))))
from sam_pth.common_tools import *


def main():
    device = device_select()
    sam2_checkpoint = "/weights/sam2.1_hiera_small.pt"
    model_cfg = "//opt/sam2/sam2/configs/sam2.1/sam2.1_hiera_s.yaml"
    # mask_generator = init_mask_generator(sam2_checkpoint, model_cfg, device)
    mask_generator = init_mask_generator_with_para(sam2_checkpoint, model_cfg, device)

    image_path = '/opt/sam2/notebooks/images/cars.jpg'
    save_dir = '/data'
    image = Image.open(image_path)
    image = np.array(image.convert("RGB"))

    plt.figure(figsize=(20, 20))
    plt.imshow(image)
    plt.axis('off')
    if show_image:
        plt.show()
    else:
        plt.savefig(osp.join(save_dir, 'cars-1.jpg'))

    masks = mask_generator.generate(image)

    '''
    掩码生成会返回一个掩码列表，其中每个掩码都是一个包含各种掩码相关数据的字典。这些键包括：
    segmentation：掩码本身；
    area：掩码的面积（以像素为单位）；
    bbox：掩码的边界框，采用 XYWH 格式；
    predicted_iou：模型自身对掩码质量的预测；
    point_coords：生成此掩码的采样输入点；
    stability_score：掩码质量的另一项度量指标；
    crop_box：用于生成此掩码的图像裁剪部分，采用 XYWH 格式。
    '''
    print(len(masks))
    print(masks[0].keys())

    plt.figure(figsize=(20, 20))
    plt.imshow(image)
    show_anns(masks)
    plt.axis('off')
    if show_image:
        plt.show()
    else:
        plt.savefig(osp.join(save_dir, 'cars-2.jpg'))

    '''
    在自动掩码生成过程中有几个可调参数，它们控制着点的采样密度以及去除低质量或重复掩码的阈值。
    此外，可以对图像的裁剪部分自动进行生成操作，以便在处理较小对象时获得更好的性能，并且后处理可以去除杂散像素和孔洞。
    以下是一个采样更多掩码的示例配置：
    '''
    mask_generator_2 = init_mask_generator_with_para(sam2_checkpoint, model_cfg, device)

    masks2 = mask_generator_2.generate(image)

    plt.figure(figsize=(20, 20))
    plt.imshow(image)
    show_anns(masks2)
    plt.axis('off')
    if show_image:
        plt.show()
    else:
        plt.savefig(osp.join(save_dir, 'cars-3.jpg'))


if __name__ == '__main__':
    main()
